import numpy as np
import matplotlib.pyplot as plt

# -------------------------- 基础配置与第一题复用参数 --------------------------
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False

# 1. 目标与导弹基础参数（与第一题完全一致）
target_true = np.array([0, 200, 0])  # 真目标底面圆心 (x,y,z)
r, h = 7, 10  # 真目标半径7m、高10m
target_fake = np.array([0, 0, 0])  # 假目标（原点）
missile_M1_init = np.array([20000, 0, 2000])  # M1初始位置
missile_speed = 300  # 导弹速度300m/s
drone_FY1_init = np.array([17800, 0, 1800])  # FY1初始位置

# 2. 烟幕基础参数（与第一题完全一致）
g = 9.8  # 重力加速度(m/s²)
smoke_radius = 10  # 烟幕有效半径10m
smoke_sink_speed = 3  # 烟幕下沉速度3m/s
smoke_valid_time = 20  # 烟幕起爆后20s有效
missile_arrive_time = np.linalg.norm(missile_M1_init - target_fake) / missile_speed  # 导弹命中假目标时间≈67.0s

# 3. 第二题优化变量边界（新增：需在约束范围内搜索最优值）
opt_bounds = {
    'v_u': [70, 140],  # 无人机速度范围 (m/s)
    'theta': [0, 2 * np.pi],  # 航向角范围 (rad，0为x轴正方向)
    't_drop': [0, 60],  # 投放时刻范围 (s，早于导弹命中)
    'delta_t': [0, 18.6]  # 投放-起爆间隔上限 (由云团不落地约束推导：Δt≤√[(1800-70)/4.9]≈18.6s)
}


# -------------------------- 第一题复用核心函数 --------------------------
# 1. 真目标外接长方体顶点生成（完全复用）
def get_cylinder_bounding_box_vertices(target_center, radius, height):
    x_min, x_max = target_center[0] - radius, target_center[0] + radius
    y_min, y_max = target_center[1] - radius, target_center[1] + radius
    z_min, z_max = target_center[2], target_center[2] + height
    vertices = [
        np.array([x_min, y_min, z_min]), np.array([x_max, y_min, z_min]),
        np.array([x_max, y_max, z_min]), np.array([x_min, y_max, z_min]),
        np.array([x_min, y_min, z_max]), np.array([x_max, y_min, z_max]),
        np.array([x_max, y_max, z_max]), np.array([x_min, y_max, z_max])
    ]
    return vertices


# 2. 导弹实时位置计算（完全复用）
def missile_real_time_position(missile_init_pos, t):
    if t < 0:
        return missile_init_pos
    # 导弹方向：从初始位置指向假目标的单位向量
    dir_vec = target_fake - missile_init_pos
    dir_vec_mag = np.linalg.norm(dir_vec)
    dir_unit_vec = dir_vec / dir_vec_mag if dir_vec_mag > 1e-6 else np.array([0, 0, 0])
    # 实时位置 = 初始位置 + 位移
    displacement = missile_speed * t * dir_unit_vec
    return missile_init_pos + displacement


# 3. 线段（导弹-顶点）与烟幕球相交判定（完全复用）
def is_segment_sphere_intersect(missile_pos, vertex_pos, smoke_center, smoke_radius):
    vec_AC = smoke_center - missile_pos  # 导弹到烟幕球心向量
    vec_AB = vertex_pos - missile_pos  # 导弹到真目标顶点向量
    ab_sq_mag = np.dot(vec_AB, vec_AB)

    if ab_sq_mag < 1e-12:  # 避免线段长度为0的计算误差
        return False

    # 参数k：定位线段上离球心最近的点
    k = np.dot(vec_AC, vec_AB) / ab_sq_mag

    # 计算最短距离
    if k <= 0.0:
        min_dist = np.linalg.norm(vec_AC)  # 最近点为导弹位置
    elif 0.0 < k < 1.0:
        # 最近点在线段内：叉积模长 / 线段长度
        cross_prod = np.cross(vec_AC, vec_AB)
        min_dist = np.linalg.norm(cross_prod) / np.sqrt(ab_sq_mag)
    else:
        min_dist = np.linalg.norm(smoke_center - vertex_pos)  # 最近点为顶点位置

    return min_dist <= smoke_radius + 1e-6  # 允许微小计算误差


# -------------------------- 第二题新增：动态参数计算（依赖优化变量） --------------------------
def calc_drone_position(opt_params, t):
    """
    计算无人机在t时刻的位置（依赖优化变量：速度v_u、航向角theta）
    opt_params: 优化变量字典 {'v_u':..., 'theta':..., 't_drop':..., 'delta_t':...}
    t: 当前时刻 (s)
    """
    v_u = opt_params['v_u']
    theta = opt_params['theta']
    # 无人机等高度飞行，z坐标始终为初始高度1800m
    x = drone_FY1_init[0] + v_u * np.cos(theta) * t
    y = drone_FY1_init[1] + v_u * np.sin(theta) * t
    z = drone_FY1_init[2]
    return np.array([x, y, z])


def calc_smoke_explosion_center(opt_params):
    """
    计算烟幕起爆点（依赖优化变量：t_drop、delta_t、v_u、theta）
    """
    t_drop = opt_params['t_drop']
    delta_t = opt_params['delta_t']
    # 1. 投放时刻（t_drop）无人机位置 = 烟幕投放点
    drop_pos = calc_drone_position(opt_params, t_drop)
    # 2. 烟幕弹从投放至起爆的运动：水平同无人机速度，竖直自由下落
    t_det = t_drop + delta_t  # 起爆时刻
    # 水平位置：无人机在t_det时刻的位置（烟幕弹水平速度=无人机速度）
    det_x, det_y = calc_drone_position(opt_params, t_det)[0:2]
    # 竖直位置：自由下落距离 = 0.5*g*delta_t²，起爆点z坐标 = 投放点z - 下落距离
    fall_dist = 0.5 * g * (delta_t ** 2)
    det_z = max(drop_pos[2] - fall_dist, 10)  # 确保云团不落地（最低z=10m）
    return np.array([det_x, det_y, det_z]), t_det


def calc_smoke_real_center(opt_params, t):
    """
    计算t时刻烟幕云团中心（依赖优化变量）
    """
    explosion_center, t_det = calc_smoke_explosion_center(opt_params)
    if t < t_det - 1e-6:  # 起爆前烟幕未形成
        return np.array([0, 0, 0])
    # 起爆后匀速下沉：z坐标 = 起爆点z - 下沉速度*(t - t_det)
    sink_dist = smoke_sink_speed * (t - t_det)
    real_z = max(explosion_center[2] - sink_dist, 10)  # 避免落地
    return np.array([explosion_center[0], explosion_center[1], real_z])


# -------------------------- 第二题新增：适应度函数（计算遮蔽时长） --------------------------
def calc_fitness(opt_params, time_step=0.05):
    """
    适应度函数：计算给定优化参数下的有效遮蔽时长
    opt_params: 优化变量字典
    time_step: 时间遍历步长（平衡精度与效率）
    """
    explosion_center, t_det = calc_smoke_explosion_center(opt_params)
    # 烟幕有效时间区间：[起爆时刻, 起爆时刻+20s] ∩ [0, 导弹命中时间]
    smoke_start = max(t_det, 0)
    smoke_end = min(t_det + smoke_valid_time, missile_arrive_time)

    if smoke_start >= smoke_end:  # 无有效时间
        return 0.0

    total_duration = 0.0
    prev_effective = False  # 前一时刻是否遮蔽
    target_vertices = get_cylinder_bounding_box_vertices(target_true, r, h)

    # 遍历时间区间计算遮蔽时长
    t = smoke_start
    while t <= smoke_end + 1e-6:
        # 1. 计算当前时刻各实体位置
        missile_pos = missile_real_time_position(missile_M1_init, t)
        smoke_center = calc_smoke_real_center(opt_params, t)

        # 2. 判断是否遮蔽（任一顶点视线被挡即有效）
        current_effective = False
        for vertex in target_vertices:
            if is_segment_sphere_intersect(missile_pos, vertex, smoke_center, smoke_radius):
                current_effective = True
                break

        # 3. 累加有效时长（连续遮蔽区间计算）
        if current_effective and not prev_effective:
            start_t = t  # 遮蔽区间开始
        elif not current_effective and prev_effective:
            total_duration += t - start_t  # 遮蔽区间结束，累加时长

        prev_effective = current_effective
        t += time_step

    # 处理区间末尾仍遮蔽的情况
    if prev_effective:
        total_duration += smoke_end - start_t

    return round(total_duration, 3)


# -------------------------- 第二题新增：遗传算法优化模块 --------------------------
def init_population(pop_size):
    """
    初始化种群：生成pop_size个优化变量个体（每个个体为字典）
    """
    population = []
    for _ in range(pop_size):
        # 在约束范围内随机生成优化变量
        v_u = np.random.uniform(opt_bounds['v_u'][0], opt_bounds['v_u'][1])
        theta = np.random.uniform(opt_bounds['theta'][0], opt_bounds['theta'][1])
        t_drop = np.random.uniform(opt_bounds['t_drop'][0], opt_bounds['t_drop'][1])
        delta_t = np.random.uniform(opt_bounds['delta_t'][0], opt_bounds['delta_t'][1])
        population.append({
            'v_u': v_u, 'theta': theta, 't_drop': t_drop, 'delta_t': delta_t
        })
    return population


def tournament_selection(population, fitness_list, k=3):
    """
    锦标赛选择：从种群中随机选k个个体，选适应度最高的进入父代池
    """
    selected = []
    pop_size = len(population)
    for _ in range(pop_size):
        # 随机选择k个个体的索引
        candidates_idx = np.random.choice(pop_size, k, replace=False)
        # 选择适应度最高的个体
        candidates_fitness = [fitness_list[idx] for idx in candidates_idx]
        best_idx = candidates_idx[np.argmax(candidates_fitness)]
        selected.append(population[best_idx])
    return selected


def arithmetic_crossover(parent1, parent2, pc=0.8):
    """
    算术交叉：对两个父代个体进行连续变量交叉（交叉概率pc）
    """
    if np.random.rand() < pc:
        # 生成交叉权重（0-1之间随机）
        alpha = np.random.rand()
        # 对每个优化变量执行交叉
        child1 = {
            'v_u': alpha * parent1['v_u'] + (1 - alpha) * parent2['v_u'],
            'theta': alpha * parent1['theta'] + (1 - alpha) * parent2['theta'],
            't_drop': alpha * parent1['t_drop'] + (1 - alpha) * parent2['t_drop'],
            'delta_t': alpha * parent1['delta_t'] + (1 - alpha) * parent2['delta_t']
        }
        child2 = {
            'v_u': (1 - alpha) * parent1['v_u'] + alpha * parent2['v_u'],
            'theta': (1 - alpha) * parent1['theta'] + alpha * parent2['theta'],
            't_drop': (1 - alpha) * parent1['t_drop'] + alpha * parent2['t_drop'],
            'delta_t': (1 - alpha) * parent1['delta_t'] + alpha * parent2['delta_t']
        }
        return child1, child2
    else:
        # 不交叉，直接返回父代
        return parent1.copy(), parent2.copy()


def gaussian_mutation(individual, pm=0.02):
    """
    高斯变异：对个体的每个优化变量添加高斯扰动（变异概率pm）
    """
    # 各变量的变异标准差（根据变量范围设置，确保扰动合理）
    sigma = {
        'v_u': 5,  # 速度扰动±5m/s
        'theta': 0.1,  # 航向角扰动±0.1rad（≈5.7°）
        't_drop': 2,  # 投放时刻扰动±2s
        'delta_t': 0.5  # 起爆间隔扰动±0.5s
    }
    for key in individual.keys():
        if np.random.rand() < pm:
            # 添加高斯扰动
            mutation = np.random.normal(0, sigma[key])
            individual[key] += mutation
            # 截断到变量边界（确保满足约束）
            individual[key] = max(opt_bounds[key][0], min(individual[key], opt_bounds[key][1]))
    return individual


def genetic_algorithm(pop_size=100, max_gen=200, pc=0.8, pm=0.02, elite_num=5):
    """
    遗传算法主函数：求解最优投放策略
    pop_size: 种群规模
    max_gen: 最大迭代次数
    pc: 交叉概率
    pm: 变异概率
    elite_num: 精英个体数（直接保留到下一代）
    """
    # 1. 初始化种群与适应度
    population = init_population(pop_size)
    fitness_list = [calc_fitness(ind) for ind in population]
    best_fitness = max(fitness_list)
    best_individual = population[np.argmax(fitness_list)].copy()
    no_improve_count = 0  # 连续无改进次数（用于提前收敛）

    # 2. 迭代进化
    for gen in range(max_gen):
        # 2.1 精英保留：保留适应度最高的elite_num个个体
        elite_indices = np.argsort(fitness_list)[-elite_num:]
        elite_pop = [population[idx].copy() for idx in elite_indices]

        # 2.2 选择操作：生成父代池
        parent_pool = tournament_selection(population, fitness_list)

        # 2.3 交叉操作：生成子代
        offspring = []
        for i in range(0, len(parent_pool), 2):
            if i + 1 >= len(parent_pool):
                offspring.append(parent_pool[i].copy())
                break
            parent1 = parent_pool[i]
            parent2 = parent_pool[i + 1]
            child1, child2 = arithmetic_crossover(parent1, parent2, pc)
            offspring.append(child1)
            offspring.append(child2)
        # 截断子代数量（确保与父代池一致）
        offspring = offspring[:pop_size - elite_num]

        # 2.4 变异操作：优化子代
        for i in range(len(offspring)):
            offspring[i] = gaussian_mutation(offspring[i], pm)

        # 2.5 生成新一代种群（精英+变异子代）
        population = elite_pop + offspring
        # 重新计算新一代适应度
        fitness_list = [calc_fitness(ind) for ind in population]

        # 2.6 更新最优个体
        current_best_fitness = max(fitness_list)
        current_best_ind = population[np.argmax(fitness_list)].copy()
        if current_best_fitness > best_fitness + 1e-6:
            best_fitness = current_best_fitness
            best_individual = current_best_ind
            no_improve_count = 0
        else:
            no_improve_count += 1

        # 2.7 打印迭代信息
        if (gen + 1) % 20 == 0 or gen == 0:
            print(f"迭代第{gen + 1:3d}代 | 当前最优遮蔽时长：{best_fitness:6.3f}s")

        # 2.8 提前收敛条件：连续10代无改进
        if no_improve_count >= 10:
            print(f"\n提前收敛（连续10代无改进），迭代终止于第{gen + 1}代")
            break

    return best_individual, best_fitness


# -------------------------- 第二题执行与结果输出 --------------------------
if __name__ == "__main__":
    print("=== 第二题：FY1单弹最优投放策略求解（遗传算法）===")
    # 运行遗传算法
    best_params, best_duration = genetic_algorithm(
        pop_size=100,  # 种群规模
        max_gen=200,  # 最大迭代次数
        pc=0.8,  # 交叉概率
        pm=0.02,  # 变异概率
        elite_num=5  # 精英个体数
    )

    # 结果整理（航向角转换为角度，方便阅读）
    best_theta_deg = best_params['theta'] * 180 / np.pi

    # 输出最优策略
    print("\n=== 最优投放策略结果 ===")
    print(f"1. 无人机飞行速度：{best_params['v_u']:6.2f} m/s")
    print(f"2. 无人机航向角：  {best_params['theta']:6.4f} rad（{best_theta_deg:6.1f}°）")
    print(f"3. 烟幕投放时刻：  {best_params['t_drop']:6.2f} s")
    print(f"4. 烟幕起爆间隔：  {best_params['delta_t']:6.2f} s")
    print(f"5. 最大有效遮蔽时长：{best_duration:6.3f} s")

    # 验证最优策略的烟幕起爆点（可选）
    explosion_center, t_det = calc_smoke_explosion_center(best_params)
    print(f"\n=== 最优策略辅助信息 ===")
    print(f"烟幕起爆时刻：{t_det:6.2f} s")
    print(f"烟幕起爆点坐标：({explosion_center[0]:.1f}, {explosion_center[1]:.1f}, {explosion_center[2]:.1f}) m")